Refactor KMeansPalettizer, add extendable training strategy capability by crowbat · Pull Request #60 · apple/coreai-optimization

crowbat · 2026-07-29T22:50:18Z

No description provided.

crowbat · 2026-08-04T22:04:10Z

+        raise NotImplementedError
+
+
+@TrainingStrategy.register("default")


Now that we have TrainingStrategyConfig being the main object users use when configuring PalettizationSpec, associating TrainingStrategy with a registry is not needed anymore (this was based on a previous design). I will remove the registry association.

u-simha

Mostly looks good, have left some comments.

One thing to note, I think this might be backward-incompatible for an older palettization checkpoint, since we have added / modified buffers?

u-simha · 2026-08-04T22:55:29Z

+    def _refresh_indices(self, weight: torch.Tensor) -> None:
+        """Recompute indices from the current centroids, without re-clustering."""
+        self.indices = self._assign_indices(weight, self.centroids).detach()
+        self._indices_stale = False


Should we re-assign only if the indices are stale? And have a flag to force re-assignment

self._indices_stale is meant to be the flag for forcing re-assignment which is checked in the forward pass during hard_assign. _refresh_indices is only ever called there if self._indices_stale is True, so I think this covers your concern?

u-simha · 2026-08-04T23:13:01Z

+        else:
+            orig_dtype = raw_lut.dtype
+            scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator.get_qparams()
+            lut = self._lut_fake_quantizer._fused_fake_quant_dequant(


Nit: lut -> fq_lut

u-simha · 2026-08-04T23:14:30Z

-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        weight = original_weights.cpu()
+    @property
+    def quantized_lut(self) -> torch.Tensor | None:


Minor suggestion: this can be reused in the lut function with a fq flag

This part, along with lut_quantization_scale and lut_quantization_zero_point were made into properties to try to preserve some backwards comaptibleness with the previous implementation which stored all 3 as separate buffers. Though there is still incompatibility when it comes to being able to save and load actual buffer values so perhaps we can consider deprecating them instead

u-simha · 2026-08-04T23:15:10Z

-        if self.enable_per_channel_scale:
-            weight = self._scale_by_per_channel_scale(weight)
+    @property
+    def lut_quantization_scale(self) -> torch.Tensor | None:


Can we have lut_quantization_qparams and return both the scale & zero point together? Any reason this is separated?

u-simha · 2026-08-04T23:19:40Z

+        with palettizer.training_mode():
+            prepared(torch.randn(2, 16)).sum().backward()
+
+        assert prepared.palettized.parametrizations.weight.original.grad is None
+        assert prepared.head.weight.grad is not None


Can we have a bit more extensive test that tests that we don't break gradients within and outside of the training mode context:

coreai-optimization/tests/quantization/test_fake_quantize.py

Line 228 in 7ad2df3

def test_qat_gradient_flow(dtype, granularity, qformulation, compression_target_tensor):

I saw a couple places we call .detach() which sometimes breaks the gradients if it is part of the backward pass (or can short circuit and have the gradients flowing through unintended variables - this happened in quantization at some point)

Sounds good, I'll add some coverage for this too

u-simha · 2026-08-06T01:48:32Z

+                        self._fp_to_schedule[param] = schedule
+                    break
+
+    def _resolve_schedule(self, module_name: str) -> PATSchedule | None:


Would this follow the module priority while applying the PAT schedule? I think I should use this for the QAT schedule too

This part of the code is pretty much a mirror of what we have in quantizer: https://github.com/apple/coreai-optimization/blob/main/src/coreai_opt/quantization/quantizer.py#L184

so it may inherently have the same issue as QAT

u-simha · 2026-08-06T01:49:40Z

+
+    def _apply_schedule(self) -> None:
+        for fp_module, schedule in self._fp_to_schedule.items():
+            fp_module.enable_fake_palett(schedule._compute_state(self._step_count))


Would it be better to do fp_module.apply(enable_fake_palett) that way even children modules have the schedule applied?

Given that fp_to_schedule is supposed to carry only leaf level modules mapping to schedules, it may be better to keep it as is to ensure we are setting the schedule for exactly the module we intend and no more accidentally

u-simha · 2026-08-06T01:52:20Z


+    @abstractmethod
    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Apply fake palettization to input tensor"""


I don't fully follow the reason for this logic to be moved to the downstream class. Is it because of the training strategy? I would think that would be generic across all fake palletize, and not specific to KMeans

At first, due to the removal of observer_enabled, most of this function simply got deleted and what was left seemed to get overridden anyways by much of the newly added things in kmeans_fake_palettizer, so it felt like this function wasn't lifting any weight.

But I think we can reframe some of what is in kmeans_fake_palettizer to be generic. I'll try some alternatives to see

u-simha · 2026-08-06T01:55:07Z

-                self._model.apply(_enable_observer)
+    @contextmanager
+    def training_mode(self):
+        """Context manager wrapping a training loop. Mutually exclusive with


We should check that the model is prepared, similar to what we do for quantization

u-simha · 2026-08-06T01:58:12Z

+        self._mode = "training"
+        try:
+            self._model.train()
+            self._build_fp_to_schedule()


Can we cache the fp_to_schedule since it shouldn't change across entering the training mode context?

I would say a common way to call this code would be to enter it on every batch step / epoch step and exit it while doing eval

Internally _build_fp_to_schedule() does check whether self._fp_to_schedule exists and returns that if it exists already. I can rename the function to _get_fp_to_schedule() to be less misleading

crowbat

Thanks for your comments @u-simha . I think on the backwards compatibility let's talk more about the best way to handle this since a good amount has changed.

crowbat · 2026-08-06T22:30:04Z

+    def _refresh_indices(self, weight: torch.Tensor) -> None:
+        """Recompute indices from the current centroids, without re-clustering."""
+        self.indices = self._assign_indices(weight, self.centroids).detach()
+        self._indices_stale = False


self._indices_stale is meant to be the flag for forcing re-assignment which is checked in the forward pass during hard_assign. _refresh_indices is only ever called there if self._indices_stale is True, so I think this covers your concern?

crowbat · 2026-08-06T22:32:17Z

+        self._mode = "training"
+        try:
+            self._model.train()
+            self._build_fp_to_schedule()


Internally _build_fp_to_schedule() does check whether self._fp_to_schedule exists and returns that if it exists already. I can rename the function to _get_fp_to_schedule() to be less misleading

crowbat · 2026-08-06T23:53:08Z


+    @abstractmethod
    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Apply fake palettization to input tensor"""


At first, due to the removal of observer_enabled, most of this function simply got deleted and what was left seemed to get overridden anyways by much of the newly added things in kmeans_fake_palettizer, so it felt like this function wasn't lifting any weight.

But I think we can reframe some of what is in kmeans_fake_palettizer to be generic. I'll try some alternatives to see

crowbat · 2026-08-07T00:23:23Z

-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        weight = original_weights.cpu()
+    @property
+    def quantized_lut(self) -> torch.Tensor | None:


This part, along with lut_quantization_scale and lut_quantization_zero_point were made into properties to try to preserve some backwards comaptibleness with the previous implementation which stored all 3 as separate buffers. Though there is still incompatibility when it comes to being able to save and load actual buffer values so perhaps we can consider deprecating them instead

crowbat · 2026-08-07T00:51:07Z


        self._num_workers = 1

+        self._mode: str = "idle"  # "idle" | "training" | "calibrating"


I'll define it in BaseModelCompressor but not add it to pruning or quantizers yet (that can come separately)

crowbat · 2026-08-07T01:07:36Z

+                        self._fp_to_schedule[param] = schedule
+                    break
+
+    def _resolve_schedule(self, module_name: str) -> PATSchedule | None:


This part of the code is pretty much a mirror of what we have in quantizer: https://github.com/apple/coreai-optimization/blob/main/src/coreai_opt/quantization/quantizer.py#L184

so it may inherently have the same issue as QAT

crowbat · 2026-08-07T01:14:56Z

+
+    def _apply_schedule(self) -> None:
+        for fp_module, schedule in self._fp_to_schedule.items():
+            fp_module.enable_fake_palett(schedule._compute_state(self._step_count))


Given that fp_to_schedule is supposed to carry only leaf level modules mapping to schedules, it may be better to keep it as is to ensure we are setting the schedule for exactly the module we intend and no more accidentally

crowbat · 2026-08-07T01:21:46Z

+        with palettizer.training_mode():
+            prepared(torch.randn(2, 16)).sum().backward()
+
+        assert prepared.palettized.parametrizations.weight.original.grad is None
+        assert prepared.head.weight.grad is not None


Sounds good, I'll add some coverage for this too

u-simha

Thanks for all the refactoring, this will be very helpful to potentially add newer palettization-based algorithms!

u-simha · 2026-08-11T20:21:40Z

+
+    def _centroids_from_lut(self, lut: torch.Tensor) -> torch.Tensor:
+        """Invert ``_reshape_lut_tensor`` to recover ``(num_blocks, num_clusters,
+        cluster_dim)`` centroids from a stored 4D LUT tensor.


Should we assert lut.ndim to be a 4D tensor or is this guaranteed?

It's not guaranteed, I'll add an error check

u-simha · 2026-08-11T20:28:28Z

+
+        assert palettizer.lut is not None
+        assert torch.isfinite(palettizer.lut).all()
+        assert palettizer.quantized_lut is not None


We can also assert that quantized_lut dtype is integer not float?

u-simha · 2026-08-11T20:29:55Z

+        assert palettizer.lut_quantization_scale is None
+        assert palettizer.lut_quantization_zero_point is None
+
+    def test_repeated_reads_idempotent_with_moving_average(self):


Isn't this testing an underlying piece of the quantization logic, which should anyway be part of our quantization tests?

Earlier when implementing the quantized_lut property I had mistakenly invoked qparams_calculator's forward pass to get the qparams instead of just querying get_qparams which caused the moving average to shift. So this is more checking that the properties are providing the requested info without updating internal state while doing so.

u-simha · 2026-08-11T20:32:08Z

+        for a, b in zip(first, second, strict=True):
+            assert torch.equal(a, b)
+
+    def test_lut_consistent_with_quantized_lut_and_scale(self):


This too feels like we are testing the "quantize" functionality

I am thinking a better test would rather be checking if the quantization error is within a certain threshold

That is, testing lut and quantized_lut have a small quantization error

Sounds good, I'll update this

u-simha · 2026-08-11T20:33:19Z

+            palettizer._blocks_to_cluster(weight_2d, axis=0)
+
+
+def _accelerator_device() -> str | None:


This can be a generic fixture in confest.py

u-simha · 2026-08-11T20:34:37Z

+
+
+@pytest.mark.skipif(_accelerator_device() is None, reason="requires a CUDA or MPS accelerator")
+def test_device_placement_on_accelerator():


I'm not sure if this is useful, but should be also test that the centroids / indices are on the right device even if the model is moved back-and-forth between CPU & GPU?

This is a good point actually. The way indices device is handled after this PR remains the same as how it was before, in that we force it to stay on CPU when we first compute it, but there are no guards in place to prevent it from being moved to GPU if the user changes the device of the model.

It's worth looking into better ways to manage this, but I'll consider that as a separate task from this PR.

crowbat requested review from anotheranshu, dengqiaoyu, pkmandke and u-simha July 29, 2026 22:50

crowbat assigned u-simha Jul 29, 2026

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 67ade25 to 2a9325e Compare July 29, 2026 22:56

crowbat commented Aug 4, 2026

View reviewed changes

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 2a9325e to 44fbc11 Compare August 4, 2026 22:15

u-simha reviewed Aug 6, 2026

View reviewed changes

crowbat commented Aug 7, 2026

View reviewed changes

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from b44e145 to 734573b Compare August 7, 2026 01:55

crowbat added 3 commits August 7, 2026 15:44

Refactor KMeansPalettizer, add extendable training strategy capability

025c0e1

Address review comments

1d00647

Restore compatibility with legacy palettizer checkpoints

a5d23ad

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 734573b to a5d23ad Compare August 7, 2026 23:47

u-simha approved these changes Aug 11, 2026

View reviewed changes

Address review comments

4fab50e

crowbat merged commit 6e4ddff into apple:main Aug 11, 2026
11 checks passed

		raise NotImplementedError


		@TrainingStrategy.register("default")


		self._num_workers = 1

		self._mode: str = "idle" # "idle" \| "training" \| "calibrating"

		palettizer._blocks_to_cluster(weight_2d, axis=0)


		def _accelerator_device() -> str \| None:



		@pytest.mark.skipif(_accelerator_device() is None, reason="requires a CUDA or MPS accelerator")
		def test_device_placement_on_accelerator():

Uh oh!

Conversation

crowbat commented Jul 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

u-simha left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crowbat left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

u-simha left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment